try:
    from mpi4py import MPI
except ImportError:
    print "cannot import MPI, continuing.."
from v3dsolver.fluid_elements.fluid_elems import FluidElem, Saffman
from numpy import binary_repr,array
from random import shuffle
import pdb

def get_domain(elems,findglobal=False):
    """
    Returns the domain of the form [xmin,ymin,zmin,length] based on given
    elems.
    
    findglobal: boolean argument, if specified True, would find global domain
                (considering elems of other processors too)
    """
    if findglobal==False: # only local computation of domain
        pos = array([elem.pos.asarray() for elem in elems])
        xmin = min(pos[:,0])
        ymin = min(pos[:,1])
        zmin = min(pos[:,2])
        xmax = max(pos[:,0])
        ymax = max(pos[:,1])
        zmax = max(pos[:,2])
        length = max([xmax-xmin,ymax-ymin,zmax-zmin])
        domain = [xmin,ymin,zmin,length]
    else:
        comm = MPI.COMM_WORLD        
        localdomain = get_domain(elems)
        globaldomain = []  
        alldomains = comm.gather(localdomain,root=0)
        if comm.rank==0:
            xmin,ymin,zmin,length = localdomain # arbitrary initialization
            for checkdomain in alldomains:
                check_xmin = checkdomain[0]
                check_ymin = checkdomain[1]
                check_zmin = checkdomain[2]
                check_xmax = checkdomain[0] + checkdomain[3]
                check_ymax = checkdomain[1] + checkdomain[3]
                check_zmax = checkdomain[2] + checkdomain[3]
                xmin = min(xmin,check_xmin)
                ymin = min(ymin,check_ymin)
                zmin = min(zmin,check_zmin)
                xmax = max(xmin+length,check_xmax)
                ymax = max(ymin+length,check_ymax)
                zmax = max(zmin+length,check_zmax)
                length = max(xmax-xmin,ymax-ymin,zmax-zmin)
            globaldomain = [xmin,ymin,zmin,length]
        globaldomain = comm.bcast(globaldomain,root=0)
        domain = globaldomain
    return domain

def get_key(pos,domain,level=2):
    """
    Returns key for sorting the elements such that nearby elements are 
    nearer to each other on the list.
    
    pos: position at which key is to be evaluated.
    domain: global domain
    level: level of resolution of the SFC
    """
    factor = 2**level
    xmin = domain[0]
    ymin = domain[1]
    zmin = domain[2]
    length = domain[3]
    x_int = int(factor * (pos[0]-xmin) / length)
    y_int = int(factor * (pos[1]-ymin) / length)
    z_int = int(factor * (pos[2]-zmin) / length)
    
    # making boundary points count in the blocks just beneath them
    if x_int==factor: x_int = factor - 1
    if y_int==factor: y_int = factor - 1
    if z_int==factor: z_int = factor - 1
    
    # binary representation and bits interleaving
    x_bin = binary_repr(x_int)
    y_bin = binary_repr(y_int) 
    z_bin = binary_repr(z_int) 
    bin_list = [x_bin,y_bin,z_bin]
    maxlen = len(binary_repr(factor))
    for i in range(len(bin_list)):
        if len(bin_list[i])<maxlen:
            bin_list[i] = '0'*(maxlen-len(bin_list[i])) + bin_list[i]
    key = ''
    for i in range(maxlen):
        key += bin_list[0][i] + bin_list[1][i] + bin_list[2][i]
    key = int(key)
    return key

class DistributionManager:
    """
    Class representing a particle distribution manager object. 
    
    Given a list of total elements, it generates keys for each element based 
    on which it distributes elements among different processors such that 
    particles close to each other remain in same processor as far as possible.
    """
    def __init__(self,comm,elems,level=2):
        """
        Initialization of the DistributionManager object. elems represent the list of 
        elements whose keys are to be managed.
        """
        self.comm = comm
        self.elems = elems # total elems (of all processors)
        self.domain = self.get_domain()
        self.level = level
    
    def get_domain(self):
        """
        Alias of get_domain function written outside DistributionManager class.
        """
        domain = []
        if self.comm.rank==0:
            # find global is False because elems itself contains particles from
            # all processors
            domain = get_domain(self.elems,findglobal=False)
        domain = self.comm.bcast(domain,root=0)
        return domain
    
    def sort_key(self,elem):
        """
        Returns the SFC key corresponding to elem
        """
        return get_key(elem.pos,self.domain,self.level)
    
    def get_this_proc_elems(self):
        """
        Returns elements of the current processor.
        """
        pos = array([elem.pos for elem in self.elems])
        i = self.comm.Get_rank()
        self.sorted_elems = []
        if i==0: # root processor
            self.sorted_elems = sorted(self.elems,key=self.sort_key)
        self.sorted_elems = self.comm.bcast(self.sorted_elems,root=0)
        # no. of particles per processor
        np_per_proc = int(round(len(self.sorted_elems)/float(self.comm.size)))
        try:
            start = i*np_per_proc
            end = min((i+1)*np_per_proc,len(self.sorted_elems))
            this_proc_elems = self.sorted_elems[start:end]
        except IndexError:
            this_proc_elems = []
        return this_proc_elems
    
    def get_this_proc_elems_random(self):
        """
        Returns elements of the current processor, shuffled randomly among 
        processors. This is used to check whether algorithm is at least correct 
        with any distribution of particles.
        """
        i = self.comm.rank
        elems_rand = self.elems[:]
        shuffle(elems_rand)
        np_per_proc = int(round(len(elems_rand)/float(self.comm.size)))
        try:
            start = i*np_per_proc
            end = min((i+1)*np_per_proc,len(elems_rand))
            this_proc_elems = elems_rand[start:end]
        except IndexError:
            this_proc_elems = []
        return this_proc_elems
